
from math import exp
import os, sys

os.environ['ETS_TOOLKIT'] = 'qt4'

# To be able to use PySide or PyQt4 and not run in conflicts with traits,
# we need to import QtGui and QtCore from pyface.qt
from pyface.qt import QtGui, QtCore
# Alternatively, you can bypass this line, but you need to make sure that
# the following lines are executed before the import of PyQT:
#   import sip
#   sip.setapi('QString', 2)

import numpy
from mayavi import mlab

from traits.api import HasTraits, Instance, on_trait_change,\
    Int, Dict
from traitsui.api import View, Item
from mayavi.core.ui.api import MayaviScene, MlabSceneModel,\
    SceneEditor

import matplotlib
from mayavi.tools.mlab_scene_model import MlabSceneModel
from traits.has_traits import HasTraits

matplotlib.use('Qt4Agg')
matplotlib.rcParams['backend.qt4']='PySide'
import pylab
import matplotlib.pyplot as plt
from PySide.QtGui import QMainWindow
from PySide import QtGui, QtCore
from mpl_toolkits.mplot3d import axes3d
from pylab import *
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas



class CompareDialog(QtGui.QWidget):
    def __init__(self, volsurface):
        #QtGui.QWidget.__init__(self)
        super(CompareDialog, self).__init__()
        layout = QtGui.QVBoxLayout(self)
        #layout.setMargin(0)
        #layout.setSpacing(0)
        self.visualization = Visualization()
        self.visualization.setVolSurface(volsurface)
        self.ui = self.visualization.edit_traits(parent=self, kind='subpanel').control
        layout.addWidget(self.ui)
        self.ui.setParent(self)


def cat(x, y, alpha, eta=1, purity=1):
    """ Multiphoton shrodinger cat. eta is the fidelity, alpha the number
        of photons"""
    cos = numpy.cos
    exp = numpy.exp
    return (1 + eta*(exp(-x**2 -(y-alpha)**2)  + exp(-x**2 -
                                                     (y+alpha)**2) + 2 * purity * exp(-x**2 - y**2) * cos(2* alpha * x))/(2
                                                                                                                          * (1 + exp(- alpha**2))))/2

class Visualization(HasTraits):
    scene = Instance(MlabSceneModel, ())


    def setVolSurface(self, volsurface):
        self.volsurface = volsurface

#    def __init__(self, volsurface):
#        self.volsurface = volsurface

    @on_trait_change('scene.activated')
    def update_plot(self):
        # This function is called when the view is opened. We don't
        # populate the scene when the view is not yet open, as some
        # VTK features require a GLContext.

        # We can do normal mlab calls on the embedded scene.
        #self.scene.mlab.test_points3d()

        x, y, z = self.volsurface.getDiscretizedXYZTuples()

        self.scene.mlab.figure(1, size=(500, 250), fgcolor=(1, 1, 1),
            bgcolor=(0.5, 0.5, 0.5))
        self.scene.mlab.clf()

        cat1 = cat(x, y, 1)
        cat2 = cat(x, y, 2)
        cat3 = cat(x, y, 3)

        #volsurf = self.volsurface.lookupVolatility(x, y)
        #vol = [[self.volsurface.lookupVolatility(K, T) for K in x] for T in y]
        volsurf = [[self.volsurface.lookupVolatility(K, T) for K in self.volsurface.strikes] for T in self.volsurface.terms]

        # The cats lie in a [0, 1] interval, with .5 being the assymptotique
        # value. We want to reposition this value to 0, so as to put it in the
        # center of our extents.
        cat1 -= 0.5
        cat2 -= 0.5
        cat3 -= 0.5

#        cat1_extent = (-14,-6, -4,4, 0,5)
#        surf_cat1 = self.scene.mlab.surf(x-10, y, cat1, colormap='Spectral', warp_scale=5,
#            extent=cat1_extent, vmin=-0.5, vmax=0.5)
#        self.scene.mlab.outline(surf_cat1, color=(.7, .7, .7))
#        self.scene.mlab.axes(surf_cat1, color=(.7, .7, .7), extent=cat1_extent,
#            ranges=(0,1, 0,1, 0,1), xlabel='', ylabel='',
#            zlabel='volatility',
#            x_axis_visibility=False, z_axis_visibility=False)


        cat1_extent = (-14,-6, -4,4, 0,5)
        surf_cat1 = self.scene.mlab.surf(x, y, volsurf, colormap='Spectral', warp_scale=5,
            vmin=-0.5, vmax=0.5)
        self.scene.mlab.outline(surf_cat1, color=(.7, .7, .7))
        self.scene.mlab.axes(surf_cat1, color=(.7, .7, .7),
        ranges=(0,1, 0,1, 0,1), xlabel='K/S', ylabel='Term',
        zlabel='volatility',
        x_axis_visibility=False, z_axis_visibility=False)


        self.scene.mlab.text(-18, -4, 'SPX', z=-4, width=0.08)

        cat2_extent = (-4,4, -4,4, 0,5)
        surf_cat2 = self.scene.mlab.surf(x, y, volsurf, colormap='Spectral', warp_scale=5,
            extent=cat2_extent, vmin=-0.5, vmax=0.5)
        self.scene.mlab.outline(surf_cat2, color=(0.7, .7, .7), extent=cat2_extent)

        self.scene.mlab.text(-4, -3, 'AAPL', z=-4, width=0.08)

        cat3_extent = (6,14, -4,4, 0,5)
        surf_cat3 = self.scene.mlab.surf(x, y, volsurf, colormap='Spectral', warp_scale=5,
            extent=cat3_extent, vmin=-0.5, vmax=0.5)
        self.scene.mlab.outline(surf_cat3, color=(.7, .7, .7), extent=cat3_extent)

        self.scene.mlab.text(6, -2.5, 'GOOG', z=-4, width=0.08)

        self.scene.mlab.title('Vol Surface Comparative view')

        self.scene.mlab.view(142, -72, 32)

        self.scene.mlab.show()


    # the layout of the dialog screated
    view = View(Item('scene', editor=SceneEditor(scene_class=MayaviScene),
        height=250, width=300, show_label=False),
        resizable=True # We need this to resize with the parent widget
    )
